#!/bin/bash
. path.sh
# Runs the "345M" parameter model

export CUDA_DEVICE_MAX_CONNECTIONS=1
export NCCL_LAUNCH_TIMEOUT=30
export NCCL_SOCKET_IFNAME=eth1
#export NCCL_IB_GID_INDEX=3
#export NCCL_IB_HCA=mlx5_2:1
#export NCCL_IB_SL=3
#export NCCL_CHECK_DISABLE=1
export NCCL_P2P_DISABLE=1
#export NCCL_LL_THRESHOLD=16384
#export NCCL_IB_CUDA_SUPPORT=1



GPUS_PER_NODE=8
# Change for multinode config
MASTER_ADDR=localhost
MASTER_PORT=12323
NNODES=1
NODE_RANK=0
WORLD_SIZE=$(($GPUS_PER_NODE*$NNODES))


VOCAB_FILE=/apdcephfs_qy3/share_976139/users/adrenzhou/nlp_workdir/pretrained_models/tencent/Hunyuan-7b-32k-llama2-chat-hf/bpe-bytelevel-vocab.json
MERGE_FILE=/apdcephfs_qy3/share_976139/users/adrenzhou/nlp_workdir/pretrained_models/tencent/Hunyuan-7b-32k-llama2-chat-hf/bpe-bytelevel-merges.txt
DATA_PATH=/apdcephfs_qy3/share_976139/users/adrenzhou/nlp_workdir/hunyuan_data/debug/packing_endwith_padding_text_document


MICRO_BATCH_SIZE=1
GLOBAL_BATCH_SIZE=4
TP=4
PP=1
CP=2
# require to align with weight dimensions
HIDDEN_SIZE=4096
FFN_HIDDEN_SIZE=11008
NUM_LAYERS=32
NUM_HEADS=32
SEQ_LENGTH=8192
######################################
SAVE_PATH=/apdcephfs_qy3/share_976139/users/adrenzhou/exp/exp_meetingGPT/llama_megatron_debug
HF_LLAMA_PATH=/apdcephfs_qy3/share_976139/users/adrenzhou/nlp_workdir/pretrained_models/tencent/Hunyuan-7b-32k-llama2-chat-hf


TOKENIZER_MODEL=/apdcephfs_qy3/share_976139/users/adrenzhou/nlp_workdir/pretrained_models/Llama-2-7b-chat-hf
# CHECKPOINT_PATH=/apdcephfs_qy3/share_976139/users/adrenzhou/nlp_workdir/Megatron-LM/hunyuan-7b-mega-T2P2

#CHECKPOINT_PATH=/apdcephfs_qy3/share_976139/users/adrenzhou/nlp_workdir/Megatron-LM/llama7b-megatron-TP${TP}-PP${PP}
CHECKPOINT_PATH=hunyuan7b-megatron-TP${TP}-PP${PP}-CP${CP}-TE
tensorboard_output=$SAVE_PATH/tensorboard
mkdir -p $SAVE_PATH/tensorboard

DISTRIBUTED_ARGS="
    --nproc_per_node $GPUS_PER_NODE \
    --nnodes $NNODES \
    --node_rank $NODE_RANK \
    --master_addr $MASTER_ADDR \
    --master_port $MASTER_PORT
"

GPT_ARGS="
    --tensor-model-parallel-size $TP \
    --pipeline-model-parallel-size $PP \
    --context-parallel-size $CP \
    --sequence-parallel \
    --num-layers $NUM_LAYERS \
    --hidden-size $HIDDEN_SIZE \
    --num-attention-heads $NUM_HEADS \
    --seq-length $SEQ_LENGTH \
    --max-position-embeddings $SEQ_LENGTH \
    --micro-batch-size $MICRO_BATCH_SIZE \
    --global-batch-size $GLOBAL_BATCH_SIZE \
    --lr 1e-5 \
    --train-iters 1000000 \
    --dataloader-type cyclic \
    --lr-decay-iters 100 \
    --lr-decay-style cosine \
    --min-lr 1e-6 \
    --weight-decay 1e-1 \
    --lr-warmup-fraction 0\
    --clip-grad 1.0 \
    --recompute-activations \
    --recompute-granularity selective \
    --use-distributed-optimizer \
    --no-load-optim \
    --no-load-rng \
    --use-rotary-position-embeddings \
    --normalization RMSNorm \
    --no-position-embedding \
    --no-masked-softmax-fusion \
    --transformer-impl transformer_engine \
    --ckpt-format torch \
    --attention-softmax-in-fp32 \
    --vocab-file $VOCAB_FILE \
    --merge-file $MERGE_FILE \
    --tokenizer-type GPT2BPETokenizer \
    --bf16
"


    #--use-legacy-models \

#    --tokenizer-type Llama2Tokenizer \
 #   --tokenizer-model ${TOKENIZER_MODEL} \

    #--untie-embeddings-and-output-weights \
    # --vocab-size 290943 \
    # --mock-data \
    # --data-path $DATA_PATH \
DATA_ARGS="
    --valid-data-path $DATA_PATH \
    --train-data-path $DATA_PATH
"

OUTPUT_ARGS="
    --log-interval 100 \
    --save-interval 1000000000 \
    --log-throughput \
    --eval-interval 1000000000 \
    --eval-iters 1
"

torchrun $DISTRIBUTED_ARGS pretrain_gpt.py \
    $GPT_ARGS \
    $DATA_ARGS \
    $OUTPUT_ARGS \
    --distributed-backend nccl \
    --save $SAVE_PATH 

